#include <chrono>
#include "lightglue.h"

LightGlue::LightGlue(std::string env_name, \
	const std::string& lg_onnx_path, \
                     cv::Size img1_size, \
                     cv::Size img2_size, \
                     int num_threads){
	env_name_ = env_name;
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, env_name_.c_str());
    Ort::SessionOptions sessionOptions;
    OrtCUDAProviderOptions options;
 	options.device_id = 0, 1, 2, 3;
	sessionOptions.AppendExecutionProvider_CUDA(options);
    sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    lg_session_ = std::shared_ptr<Ort::Session>(new Ort::Session(env, lg_onnx_path.c_str(), sessionOptions));
    img1_size_ = img1_size;
    img2_size_ = img2_size;
}

std::vector<cv::DMatch> LightGlue::L_Match(const std::vector<cv::KeyPoint>key_points1, \
                                           const std::vector<cv::KeyPoint>key_points2, \
                                           const cv::Mat& descriptors1, \
                                           const cv::Mat& descriptors2){
    Ort::Env env(ORT_LOGGING_LEVEL_FATAL, env_name_.c_str());
    // 1.pre processing
	// auto t1 = std::chrono::high_resolution_clock::now();
    vector<float>kp1, kp2; 
    kp1.resize(key_points1.size() * 2); 
    kp2.resize(key_points2.size() * 2);
	float f1wid = img1_size_.width / 2.0f;
	float f1hei = img1_size_.height / 2.0f;
    for (int i = 0; i < key_points1.size(); i++){
		kp1[2 * i] = (key_points1[i].pt.x - f1wid) / f1wid;
		kp1[2 * i + 1] = (key_points1[i].pt.y - f1hei) / f1hei;
	}

    float f2wid = img2_size_.width / 2.0f;
	float f2hei = img2_size_.height / 2.0f;
	for (int i = 0; i < key_points2.size(); i++)
	{
		kp2[2 * i] = (key_points2[i].pt.x - f2wid) / f2wid;
		kp2[2 * i + 1] = (key_points2[i].pt.y - f2hei) / f2hei;
	}

    vector<float>des1; des1.resize(key_points1.size() * 256);
	Mat des1mat = descriptors1;
	for (int w = 0; w < des1mat.cols; w++)
	{
		for (int h = 0; h < des1mat.rows; h++)
		{
			int index = h * descriptors1.cols + w;
			des1[index] = des1mat.at<float>(h, w);
		}
	}
	vector<float>des2; des2.resize(key_points2.size() * 256);
	Mat des2mat = descriptors2;
	for (int w = 0; w < des2mat.cols; w++)
	{
		for (int h = 0; h < des2mat.rows; h++)
		{
			int index = h * descriptors2.cols + w;
			des2[index] = des2mat.at<float>(h, w);
		}
	}
	// auto t2 = std::chrono::high_resolution_clock::now();
	// auto t12 = std::chrono::duration_cast<std::chrono::microseconds>(t2-t1);
	// std::cout<<"Pre time is: "<<t12.count()*1.0/1000<<" ms"<<std::endl;

    // 2.infer
    const char* input_names[] = { "kpts0", "kpts1", "desc0", "desc1" };
	Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    std::vector<Ort::Value> inputTensor;
	vector<int64_t> kp1Shape{ 1,(int64_t)key_points1.size(), 2 };
	inputTensor.emplace_back(Ort::Value::CreateTensor<float>(memoryInfo, kp1.data(), kp1.size(), kp1Shape.data(), kp1Shape.size()));
	vector<int64_t> kp2Shape{ 1,(int64_t)key_points2.size(), 2 };
	inputTensor.emplace_back(Ort::Value::CreateTensor<float>(memoryInfo, kp2.data(), kp2.size(), kp2Shape.data(), kp2Shape.size()));
	vector<int64_t> des1Shape{ 1,(int64_t)key_points1.size(), descriptors1.cols };
	inputTensor.emplace_back(Ort::Value::CreateTensor<float>(memoryInfo, des1.data(), des1.size(), des1Shape.data(), des1Shape.size()));
	vector<int64_t> des2Shape{ 1,(int64_t)key_points2.size(), descriptors2.cols };
	inputTensor.emplace_back(Ort::Value::CreateTensor<float>(memoryInfo, des2.data(), des2.size(), des2Shape.data(), des2Shape.size()));
	const char* output_names[] = { "matches0","matches1","mscores0","mscores1" };
	Ort::RunOptions run_options;
	vector<Ort::Value> outputs = lg_session_->Run(run_options, input_names, inputTensor.data(), 4, output_names, 4);
	// auto t3 = std::chrono::high_resolution_clock::now();
	// auto t23 = std::chrono::duration_cast<std::chrono::microseconds>(t3-t2);
	// std::cout<<"Infer time is: "<<t23.count()*1.0/1000<<" ms"<<std::endl;

    // 3.parse output
    std::vector<int64_t> match1shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
	int64_t* match1 = (int64_t*)outputs[0].GetTensorMutableData<void>();
	int match1counts = match1shape[1];

	std::vector<int64_t> mscoreshape = outputs[2].GetTensorTypeAndShapeInfo().GetShape();
	float* mscore1 = (float*)outputs[2].GetTensorMutableData<void>();
	int mscorecounts = mscoreshape[1];

	std::vector<int64_t> match2shape = outputs[1].GetTensorTypeAndShapeInfo().GetShape();
	int64_t* match2 = (int64_t*)outputs[1].GetTensorMutableData<void>();
	int match2counts = match2shape[1];

    std::vector<cv::DMatch> matches;
	for (int i = 0; i < match1counts; i++)
	{
		// std::cout<<"mscore1=================== " <<mscore1[i]<<std::endl;
		if (match1[i] > -1 && mscore1[i] > 0.0f && match2[match1[i]] == i)
		{
			cv::DMatch mt;
			mt.queryIdx = i;
			mt.trainIdx = match1[i];
			matches.push_back(mt);
		}
	}
	// auto t4 = std::chrono::high_resolution_clock::now();
	// auto t34 = std::chrono::duration_cast<std::chrono::microseconds>(t4-t3);
	// std::cout<<"After time is: "<<t34.count()*1.0/1000<<" ms"<<std::endl;
	// std::cout << "matches count:" << matches.size() << std::endl;
    return matches;
}